package com.mcp.mslocalmcp.controller;

import com.mcp.mslocalmcp.model.*;
import com.mcp.mslocalmcp.service.MySqlService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;

@RestController
@RequestMapping("/mcp")
@CrossOrigin(origins = "*")
public class McpController {
    
    private static final Logger logger = LoggerFactory.getLogger(McpController.class);
    
    @Autowired
    private MySqlService mySqlService;
    
    @PostMapping("/call")
    public ResponseEntity<McpResponse> handleMcpCall(@RequestBody McpRequest request) {
        logger.info("收到MCP请求: {}", request.getMethod());
        
        try {
            String method = request.getMethod();
            String name = request.getParams() != null ? request.getParams().getName() : null;
            Map<String, Object> arguments = request.getParams() != null ? request.getParams().getArguments() : null;
            
            Object result = null;
            
            switch (method) {
                case "tools/call":
                    result = handleToolCall(name, arguments);
                    break;
                case "initialize":
                    result = handleInitialize(arguments);
                    break;
                default:
                    return ResponseEntity.ok(new McpResponse("2.0", request.getId(), 
                        new McpError(-32601, "方法不存在: " + method)));
            }
            
            return ResponseEntity.ok(new McpResponse("2.0", request.getId(), result));
            
        } catch (Exception e) {
            logger.error("处理MCP请求时发生错误: {}", e.getMessage(), e);
            return ResponseEntity.ok(new McpResponse("2.0", request.getId(), 
                new McpError(-32603, "内部错误: " + e.getMessage())));
        }
    }
    
    @GetMapping(value = "/call", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public ResponseEntity<String> handleMcpStream() {
        logger.info("收到MCP流式连接请求");
        
        // 返回流式响应，包含工具列表
        StringBuilder response = new StringBuilder();
        response.append("data: {\n");
        response.append("  \"jsonrpc\": \"2.0\",\n");
        response.append("  \"method\": \"tools/list\",\n");
        response.append("  \"params\": {\n");
        response.append("    \"tools\": [\n");
        response.append("      {\n");
        response.append("        \"name\": \"test_connection\",\n");
        response.append("        \"description\": \"测试数据库连接\",\n");
        response.append("        \"inputSchema\": {\n");
        response.append("          \"type\": \"object\",\n");
        response.append("          \"properties\": {}\n");
        response.append("        }\n");
        response.append("      },\n");
        response.append("      {\n");
        response.append("        \"name\": \"query\",\n");
        response.append("        \"description\": \"执行SQL查询\",\n");
        response.append("        \"inputSchema\": {\n");
        response.append("          \"type\": \"object\",\n");
        response.append("          \"properties\": {\n");
        response.append("            \"sql\": {\n");
        response.append("              \"type\": \"string\",\n");
        response.append("              \"description\": \"SQL查询语句\"\n");
        response.append("            },\n");
        response.append("            \"params\": {\n");
        response.append("              \"type\": \"array\",\n");
        response.append("              \"description\": \"查询参数\"\n");
        response.append("            }\n");
        response.append("          },\n");
        response.append("          \"required\": [\"sql\"]\n");
        response.append("        }\n");
        response.append("      },\n");
        response.append("      {\n");
        response.append("        \"name\": \"execute\",\n");
        response.append("        \"description\": \"执行SQL更新操作\",\n");
        response.append("        \"inputSchema\": {\n");
        response.append("          \"type\": \"object\",\n");
        response.append("          \"properties\": {\n");
        response.append("            \"sql\": {\n");
        response.append("              \"type\": \"string\",\n");
        response.append("              \"description\": \"SQL语句\"\n");
        response.append("            },\n");
        response.append("            \"params\": {\n");
        response.append("              \"type\": \"array\",\n");
        response.append("              \"description\": \"参数\"\n");
        response.append("            }\n");
        response.append("          },\n");
        response.append("          \"required\": [\"sql\"]\n");
        response.append("        }\n");
        response.append("      },\n");
        response.append("      {\n");
        response.append("        \"name\": \"list_tables\",\n");
        response.append("        \"description\": \"获取数据库中的所有表\",\n");
        response.append("        \"inputSchema\": {\n");
        response.append("          \"type\": \"object\",\n");
        response.append("          \"properties\": {}\n");
        response.append("        }\n");
        response.append("      },\n");
        response.append("      {\n");
        response.append("        \"name\": \"describe_table\",\n");
        response.append("        \"description\": \"获取表结构\",\n");
        response.append("        \"inputSchema\": {\n");
        response.append("          \"type\": \"object\",\n");
        response.append("          \"properties\": {\n");
        response.append("            \"table\": {\n");
        response.append("              \"type\": \"string\",\n");
        response.append("              \"description\": \"表名\"\n");
        response.append("            }\n");
        response.append("          },\n");
        response.append("          \"required\": [\"table\"]\n");
        response.append("        }\n");
        response.append("      }\n");
        response.append("    ]\n");
        response.append("  }\n");
        response.append("}\n\n");
        
        return ResponseEntity.ok()
            .contentType(MediaType.TEXT_EVENT_STREAM)
            .body(response.toString());
    }
    
    private Object handleInitialize(Map<String, Object> arguments) {
        Map<String, Object> result = new HashMap<>();
        result.put("protocolVersion", "2024-11-05");
        result.put("capabilities", getCapabilities());
        result.put("serverInfo", getServerInfo());
        return result;
    }
    
    private Object handleToolCall(String name, Map<String, Object> arguments) {
        if (arguments == null) {
            arguments = new HashMap<>();
        }
        
        switch (name) {
            case "connect_db":
                return handleConnectDb(arguments);
            case "query":
                return handleQuery(arguments);
            case "execute":
                return handleExecute(arguments);
            case "list_tables":
                return handleListTables();
            case "describe_table":
                return handleDescribeTable(arguments);
            case "test_connection":
                return handleTestConnection();
            default:
                throw new RuntimeException("未知的工具: " + name);
        }
    }
    
    private Object handleConnectDb(Map<String, Object> arguments) {
        // 连接数据库的逻辑已经在配置中处理
        return Map.of("status", "connected", "message", "数据库连接已建立");
    }
    
    private Object handleQuery(Map<String, Object> arguments) {
        String sql = (String) arguments.get("sql");
        @SuppressWarnings("unchecked")
        List<Object> params = (List<Object>) arguments.get("params");
        
        if (sql == null) {
            throw new RuntimeException("SQL查询不能为空");
        }
        
        Object[] paramArray = params != null ? params.toArray() : new Object[0];
        List<Map<String, Object>> result = mySqlService.executeQuery(sql, paramArray);
        
        return Map.of(
            "rows", result,
            "count", result.size()
        );
    }
    
    private Object handleExecute(Map<String, Object> arguments) {
        String sql = (String) arguments.get("sql");
        @SuppressWarnings("unchecked")
        List<Object> params = (List<Object>) arguments.get("params");
        
        if (sql == null) {
            throw new RuntimeException("SQL语句不能为空");
        }
        
        Object[] paramArray = params != null ? params.toArray() : new Object[0];
        int affectedRows = mySqlService.executeUpdate(sql, paramArray);
        
        return Map.of(
            "affectedRows", affectedRows,
            "message", "操作执行成功"
        );
    }
    
    private Object handleListTables() {
        List<String> tables = mySqlService.listTables();
        return Map.of(
            "tables", tables,
            "count", tables.size()
        );
    }
    
    private Object handleDescribeTable(Map<String, Object> arguments) {
        String tableName = (String) arguments.get("table");
        
        if (tableName == null) {
            throw new RuntimeException("表名不能为空");
        }
        
        List<Map<String, Object>> structure = mySqlService.describeTable(tableName);
        return Map.of(
            "table", tableName,
            "structure", structure
        );
    }
    
    private Object handleTestConnection() {
        boolean connected = mySqlService.testConnection();
        return Map.of(
            "connected", connected,
            "message", connected ? "连接正常" : "连接失败"
        );
    }
    
    private Map<String, Object> getCapabilities() {
        Map<String, Object> capabilities = new HashMap<>();
        capabilities.put("tools", getTools());
        return capabilities;
    }
    
    private List<Map<String, Object>> getTools() {
        List<Map<String, Object>> tools = new ArrayList<>();
        
        // 连接数据库工具
        tools.add(Map.of(
            "name", "connect_db",
            "description", "连接到MySQL数据库",
            "inputSchema", Map.of(
                "type", "object",
                "properties", Map.of(
                    "host", Map.of("type", "string", "description", "数据库主机"),
                    "user", Map.of("type", "string", "description", "用户名"),
                    "password", Map.of("type", "string", "description", "密码"),
                    "database", Map.of("type", "string", "description", "数据库名")
                )
            )
        ));
        
        // 查询工具
        tools.add(Map.of(
            "name", "query",
            "description", "执行SELECT查询",
            "inputSchema", Map.of(
                "type", "object",
                "properties", Map.of(
                    "sql", Map.of("type", "string", "description", "SQL查询语句"),
                    "params", Map.of("type", "array", "description", "查询参数")
                ),
                "required", List.of("sql")
            )
        ));
        
        // 执行工具
        tools.add(Map.of(
            "name", "execute",
            "description", "执行INSERT、UPDATE、DELETE操作",
            "inputSchema", Map.of(
                "type", "object",
                "properties", Map.of(
                    "sql", Map.of("type", "string", "description", "SQL语句"),
                    "params", Map.of("type", "array", "description", "参数")
                ),
                "required", List.of("sql")
            )
        ));
        
        // 列出表工具
        tools.add(Map.of(
            "name", "list_tables",
            "description", "获取数据库中的所有表"
        ));
        
        // 描述表工具
        tools.add(Map.of(
            "name", "describe_table",
            "description", "获取表结构",
            "inputSchema", Map.of(
                "type", "object",
                "properties", Map.of(
                    "table", Map.of("type", "string", "description", "表名")
                ),
                "required", List.of("table")
            )
        ));
        
        // 测试连接工具
        tools.add(Map.of(
            "name", "test_connection",
            "description", "测试数据库连接"
        ));
        
        return tools;
    }
    
    private Map<String, Object> getServerInfo() {
        return Map.of(
            "name", "mysql-mcp-server",
            "version", "1.0.0"
        );
    }
} 